{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from comet_ml import Experiment\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "from fastai.vision import *\n",
    "import torch\n",
    "from torchsummary import summary\n",
    "torch.cuda.set_device(0)\n",
    "torch.manual_seed(0)\n",
    "torch.cuda.manual_seed(0)\n",
    "\n",
    "# stage should be in 0 to 5 (5 for classifier stage)\n",
    "hyper_params = {\n",
    "    \"stage\": 0,\n",
    "    \"repeated\": 0,\n",
    "    \"num_classes\": 10,\n",
    "    \"batch_size\": 64,\n",
    "    \"num_epochs\": 100,\n",
    "    \"learning_rate\": 1e-4\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = untar_data(URLs.IMAGENETTE)\n",
    "new_path = path/'new'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "tfms = get_transforms(do_flip=False)\n",
    "data = ImageDataBunch.from_folder(new_path, train = 'train', valid = 'val', test = 'test', bs = hyper_params[\"batch_size\"], size = 224, ds_tfms = tfms).normalize(imagenet_stats)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data.show_batch(rows = 8, ds_type = 'val')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model on GPU\n",
      "0.0.weight torch.Size([64, 3, 7, 7])\n",
      "True\n",
      "0.2.weight torch.Size([64])\n",
      "True\n",
      "0.2.bias torch.Size([64])\n",
      "True\n",
      "2.0.0.weight torch.Size([64, 64, 3, 3])\n",
      "False\n",
      "2.0.2.weight torch.Size([64])\n",
      "False\n",
      "2.0.2.bias torch.Size([64])\n",
      "False\n",
      "2.1.conv1.0.weight torch.Size([64, 64, 3, 3])\n",
      "False\n",
      "2.1.conv1.2.weight torch.Size([64])\n",
      "False\n",
      "2.1.conv1.2.bias torch.Size([64])\n",
      "False\n",
      "3.0.0.weight torch.Size([128, 64, 3, 3])\n",
      "False\n",
      "3.0.2.weight torch.Size([128])\n",
      "False\n",
      "3.0.2.bias torch.Size([128])\n",
      "False\n",
      "3.1.conv1.0.weight torch.Size([128, 128, 3, 3])\n",
      "False\n",
      "3.1.conv1.2.weight torch.Size([128])\n",
      "False\n",
      "3.1.conv1.2.bias torch.Size([128])\n",
      "False\n",
      "4.0.0.weight torch.Size([256, 128, 3, 3])\n",
      "False\n",
      "4.0.2.weight torch.Size([256])\n",
      "False\n",
      "4.0.2.bias torch.Size([256])\n",
      "False\n",
      "4.1.conv1.0.weight torch.Size([256, 256, 3, 3])\n",
      "False\n",
      "4.1.conv1.2.weight torch.Size([256])\n",
      "False\n",
      "4.1.conv1.2.bias torch.Size([256])\n",
      "False\n",
      "5.0.0.weight torch.Size([512, 256, 3, 3])\n",
      "False\n",
      "5.0.2.weight torch.Size([512])\n",
      "False\n",
      "5.0.2.bias torch.Size([512])\n",
      "False\n",
      "5.1.conv1.0.weight torch.Size([512, 512, 3, 3])\n",
      "False\n",
      "5.1.conv1.2.weight torch.Size([512])\n",
      "False\n",
      "5.1.conv1.2.bias torch.Size([512])\n",
      "False\n",
      "8.weight torch.Size([256, 1024])\n",
      "False\n",
      "8.bias torch.Size([256])\n",
      "False\n",
      "9.weight torch.Size([10, 256])\n",
      "False\n",
      "9.bias torch.Size([10])\n",
      "False\n"
     ]
    }
   ],
   "source": [
    "learn = cnn_learner(data, models.resnet34, metrics = accuracy)\n",
    "learn = learn.load('/home/akshay/.fastai/data/cifar/models/resnet34_cifar_bs64')\n",
    "learn.freeze()\n",
    "# learn.summary()\n",
    "\n",
    "class Flatten(nn.Module) :\n",
    "    def forward(self, input):\n",
    "        return input.view(input.size(0), -1)\n",
    "\n",
    "def conv2(ni, nf) : \n",
    "    return conv_layer(ni, nf, stride = 2)\n",
    "\n",
    "class ResBlock(nn.Module):\n",
    "    def __init__(self, nf):\n",
    "        super().__init__()\n",
    "        self.conv1 = conv_layer(nf,nf)\n",
    "        \n",
    "    def forward(self, x): \n",
    "        return (x + self.conv1(x))\n",
    "\n",
    "def conv_and_res(ni, nf): \n",
    "    return nn.Sequential(conv2(ni, nf), ResBlock(nf))\n",
    "\n",
    "def conv_(nf) : \n",
    "    return nn.Sequential(conv_layer(nf, nf), ResBlock(nf))\n",
    "    \n",
    "net = nn.Sequential(\n",
    "    conv_layer(3, 64, ks = 7, stride = 2, padding = 3),\n",
    "    nn.MaxPool2d(3, 2, padding = 1),\n",
    "    conv_(64),\n",
    "    conv_and_res(64, 128),\n",
    "    conv_and_res(128, 256),\n",
    "    conv_and_res(256, 512),\n",
    "    AdaptiveConcatPool2d(),\n",
    "    Flatten(),\n",
    "    nn.Linear(2 * 512, 256),\n",
    "    nn.Linear(256, 10)\n",
    ")\n",
    "\n",
    "net.cpu()\n",
    "if hyper_params['stage'] != 0 : \n",
    "    filename = '../saved_models/smallest_stage' + str(hyper_params['stage']) + '/model' + str(hyper_params['repeated']) + '.pt'\n",
    "    net.load_state_dict(torch.load(filename, map_location = 'cpu'))\n",
    "\n",
    "if torch.cuda.is_available() : \n",
    "    net = net.cuda()\n",
    "    print('Model on GPU')\n",
    "    \n",
    "for name, param in net.named_parameters() : \n",
    "    print(name, param.shape)\n",
    "    param.requires_grad = False\n",
    "    if name[0] == str(hyper_params['stage'] + 1) and hyper_params['stage'] != 0 :\n",
    "        param.requires_grad = True\n",
    "    elif name[0] == str(hyper_params['stage']) and hyper_params['stage'] == 0 : \n",
    "        param.requires_grad = True\n",
    "    print(param.requires_grad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1         [-1, 64, 112, 112]           9,408\n",
      "              ReLU-2         [-1, 64, 112, 112]               0\n",
      "       BatchNorm2d-3         [-1, 64, 112, 112]             128\n",
      "         MaxPool2d-4           [-1, 64, 56, 56]               0\n",
      "            Conv2d-5          [-1, 128, 28, 28]          73,728\n",
      "              ReLU-6          [-1, 128, 28, 28]               0\n",
      "       BatchNorm2d-7          [-1, 128, 28, 28]             256\n",
      "            Conv2d-8          [-1, 256, 14, 14]         294,912\n",
      "              ReLU-9          [-1, 256, 14, 14]               0\n",
      "      BatchNorm2d-10          [-1, 256, 14, 14]             512\n",
      "           Conv2d-11            [-1, 512, 7, 7]       1,179,648\n",
      "             ReLU-12            [-1, 512, 7, 7]               0\n",
      "      BatchNorm2d-13            [-1, 512, 7, 7]           1,024\n",
      "AdaptiveMaxPool2d-14            [-1, 512, 1, 1]               0\n",
      "AdaptiveAvgPool2d-15            [-1, 512, 1, 1]               0\n",
      "AdaptiveConcatPool2d-16           [-1, 1024, 1, 1]               0\n",
      "          Flatten-17                 [-1, 1024]               0\n",
      "           Linear-18                  [-1, 256]         262,400\n",
      "           Linear-19                   [-1, 10]           2,570\n",
      "================================================================\n",
      "Total params: 1,824,586\n",
      "Trainable params: 9,536\n",
      "Non-trainable params: 1,815,050\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 0.57\n",
      "Forward/backward pass size (MB): 23.95\n",
      "Params size (MB): 6.96\n",
      "Estimated Total Size (MB): 31.49\n",
      "----------------------------------------------------------------\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (0): Sequential(\n",
       "    (0): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
       "    (1): ReLU(inplace)\n",
       "    (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (1): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
       "  (2): Sequential(\n",
       "    (0): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "    (1): ReLU(inplace)\n",
       "    (2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (3): Sequential(\n",
       "    (0): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "    (1): ReLU(inplace)\n",
       "    (2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (4): Sequential(\n",
       "    (0): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "    (1): ReLU(inplace)\n",
       "    (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (5): AdaptiveConcatPool2d(\n",
       "    (ap): AdaptiveAvgPool2d(output_size=1)\n",
       "    (mp): AdaptiveMaxPool2d(output_size=1)\n",
       "  )\n",
       "  (6): Flatten()\n",
       "  (7): Linear(in_features=1024, out_features=256, bias=True)\n",
       "  (8): Linear(in_features=256, out_features=10, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# x, y = next(iter(data.train_dl))\n",
    "# net(torch.autograd.Variable(x).cuda())\n",
    "summary(net, (3, 224, 224))\n",
    "# print(learn.summary())\n",
    "net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SaveFeatures :\n",
    "    def __init__(self, m) : \n",
    "        self.handle = m.register_forward_hook(self.hook_fn)\n",
    "    def hook_fn(self, m, inp, outp) : \n",
    "        self.features = outp\n",
    "    def remove(self) :\n",
    "        self.handle.remove()\n",
    "        \n",
    "# saving outputs of all Basic Blocks\n",
    "mdl = learn.model\n",
    "sf = [SaveFeatures(m) for m in [mdl[0][2], mdl[0][5], mdl[0][6], mdl[0][7]]]\n",
    "sf2 = [SaveFeatures(m) for m in [net[0], net[2], net[3], net[4]]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([64, 64, 112, 112])\n",
      "torch.Size([64, 128, 28, 28])\n",
      "torch.Size([64, 256, 14, 14])\n",
      "torch.Size([64, 512, 7, 7])\n"
     ]
    }
   ],
   "source": [
    "x, y = next(iter(data.train_dl))\n",
    "x = torch.autograd.Variable(x).cuda()\n",
    "out1 = mdl(x)\n",
    "out2 = net(x)\n",
    "for i in range(4) : \n",
    "    print(sf[i].features.shape)\n",
    "    assert(sf[i].features.shape == sf2[i].features.shape)\n",
    "del x, y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Stage-wise training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if hyper_params['stage'] == 0 : \n",
    "    filename = '../saved_models/smallest_stage' + str(hyper_params['stage']) + '/model' + str(hyper_params['repeated']) + '.pt'\n",
    "else : \n",
    "    filename = '../saved_models/smallest_stage' + str(hyper_params['stage'] + 1) + '/model' + str(hyper_params['repeated']) + '.pt'\n",
    "optimizer = torch.optim.Adam(net.parameters(), lr = hyper_params[\"learning_rate\"])\n",
    "total_step = len(data.train_ds) // hyper_params[\"batch_size\"]\n",
    "train_loss_list = list()\n",
    "val_loss_list = list()\n",
    "min_val = 100\n",
    "for epoch in range(hyper_params[\"num_epochs\"]):\n",
    "    trn = []\n",
    "    net.train()\n",
    "    for i, (images, labels) in enumerate(data.train_dl) :\n",
    "        if torch.cuda.is_available():\n",
    "            images = torch.autograd.Variable(images).cuda().float()\n",
    "            labels = torch.autograd.Variable(labels).cuda()\n",
    "        else : \n",
    "            images = torch.autograd.Variable(images).float()\n",
    "            labels = torch.autograd.Variable(labels)\n",
    "\n",
    "        y_pred = net(images)\n",
    "        y_pred2 = mdl(images)\n",
    "\n",
    "        loss = F.mse_loss(sf2[hyper_params[\"stage\"]].features, sf[hyper_params[\"stage\"]].features)\n",
    "        trn.append(loss.item())\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "#         torch.nn.utils.clip_grad_value_(net.parameters(), 10)\n",
    "        optimizer.step()\n",
    "\n",
    "        if i % 50 == 49 :\n",
    "            print('epoch = ', epoch + 1, ' step = ', i + 1, ' of total steps ', total_step, ' loss = ', loss.item())\n",
    "\n",
    "    train_loss = (sum(trn) / len(trn))\n",
    "    train_loss_list.append(train_loss)\n",
    "\n",
    "    net.eval()\n",
    "    val = []\n",
    "    with torch.no_grad() :\n",
    "        for i, (images, labels) in enumerate(data.valid_dl) :\n",
    "            if torch.cuda.is_available():\n",
    "                images = torch.autograd.Variable(images).cuda().float()\n",
    "                labels = torch.autograd.Variable(labels).cuda()\n",
    "            else : \n",
    "                images = torch.autograd.Variable(images).float()\n",
    "                labels = torch.autograd.Variable(labels)\n",
    "\n",
    "            # Forward pass\n",
    "            y_pred = net(images)\n",
    "            y_pred2 = mdl(images)\n",
    "            loss = F.mse_loss(sf[hyper_params[\"stage\"]].features, sf2[hyper_params[\"stage\"]].features)\n",
    "            val.append(loss.item())\n",
    "\n",
    "    val_loss = sum(val) / len(val)\n",
    "    val_loss_list.append(val_loss)\n",
    "    print('epoch : ', epoch + 1, ' / ', hyper_params[\"num_epochs\"], ' | TL : ', train_loss, ' | VL : ', val_loss)\n",
    "    experiment.log_metric(\"train_loss\", train_loss)\n",
    "    experiment.log_metric(\"val_loss\", val_loss)\n",
    "\n",
    "    if val_loss < min_val :\n",
    "        print('saving model')\n",
    "        min_val = val_loss\n",
    "        torch.save(net.state_dict(), filename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# learn = Learner(data, net, metrics = accuracy)\n",
    "net.cpu()\n",
    "net.load_state_dict(torch.load('../saved_models/stage5/model0.pt', map_location = 'cpu'))\n",
    "net = net.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(range(hyper_params[\"num_epochs\"]), train_loss_list, label = 'train_loss') \n",
    "plt.plot(range(hyper_params[\"num_epochs\"]), val_loss_list, label = 'val_loss')\n",
    "plt.legend()\n",
    "plt.savefig('../figures/stage5/train_loss.jpg')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Training the classifier only after stage-wise training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "hyper_params = {\n",
    "    \"stage\": 5,\n",
    "    \"repeated\": 3,\n",
    "    \"num_classes\": 10,\n",
    "    \"batch_size\": 64,\n",
    "    \"num_epochs\": 100,\n",
    "    \"learning_rate\": 1e-4\n",
    "}\n",
    "\n",
    "class Flatten(nn.Module) :\n",
    "    def forward(self, input):\n",
    "        return input.view(input.size(0), -1)\n",
    "\n",
    "def conv2(ni, nf) : \n",
    "    return conv_layer(ni, nf, stride = 2)\n",
    "\n",
    "class ResBlock(nn.Module):\n",
    "    def __init__(self, nf):\n",
    "        super().__init__()\n",
    "        self.conv1 = conv_layer(nf,nf)\n",
    "        \n",
    "    def forward(self, x): \n",
    "        return (x + self.conv1(x))\n",
    "\n",
    "def conv_and_res(ni, nf): \n",
    "    return nn.Sequential(conv2(ni, nf), ResBlock(nf))\n",
    "\n",
    "def conv_(nf) : \n",
    "    return nn.Sequential(conv_layer(nf, nf), ResBlock(nf))\n",
    "    \n",
    "net = nn.Sequential(\n",
    "    conv_layer(3, 64, ks = 7, stride = 2, padding = 3),\n",
    "    nn.MaxPool2d(3, 2, padding = 1),\n",
    "    conv_(64),\n",
    "    conv_and_res(64, 128),\n",
    "    conv_and_res(128, 256),\n",
    "    AdaptiveConcatPool2d(),\n",
    "    Flatten(),\n",
    "    nn.Linear(2 * 256, 128),\n",
    "    nn.Linear(128, hyper_params[\"num_classes\"])\n",
    ")\n",
    "\n",
    "net.cpu()\n",
    "net.load_state_dict(torch.load('../saved_models/small_stage4/model1.pt', map_location = 'cpu'))\n",
    "\n",
    "if torch.cuda.is_available() : \n",
    "    net = net.cuda()\n",
    "    print('Model on GPU')\n",
    "    \n",
    "for name, param in net.named_parameters() : \n",
    "    print(name, param.shape)\n",
    "    param.requires_grad = False\n",
    "    if name[0] == '7' or name[0] == '8':\n",
    "        param.requires_grad = True\n",
    "    print(param.requires_grad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _get_accuracy(dataloader, Net):\n",
    "    total = 0\n",
    "    correct = 0\n",
    "    Net.eval()\n",
    "    for i, (images, labels) in enumerate(dataloader):\n",
    "        images = torch.autograd.Variable(images).float()\n",
    "        labels = torch.autograd.Variable(labels).float()\n",
    "        \n",
    "        if torch.cuda.is_available() : \n",
    "            images = images.cuda()\n",
    "            labels = labels.cuda()\n",
    "\n",
    "        outputs = Net.forward(images)\n",
    "        outputs = F.log_softmax(outputs, dim = 1)\n",
    "\n",
    "        _, pred_ind = torch.max(outputs, 1)\n",
    "        \n",
    "        # converting to numpy arrays\n",
    "        labels = labels.data.cpu().numpy()\n",
    "        pred_ind = pred_ind.data.cpu().numpy()\n",
    "        \n",
    "        # get difference\n",
    "        diff_ind = labels - pred_ind\n",
    "        # correctly classified will be 1 and will get added\n",
    "        # incorrectly classified will be 0\n",
    "        correct += np.count_nonzero(diff_ind == 0)\n",
    "        total += len(diff_ind)\n",
    "\n",
    "    accuracy = correct / total\n",
    "    # print(len(diff_ind))\n",
    "    return accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(net.parameters(), lr = hyper_params[\"learning_rate\"])\n",
    "total_step = len(data.train_ds) // hyper_params[\"batch_size\"]\n",
    "train_loss_list = list()\n",
    "val_loss_list = list()\n",
    "min_val = 0\n",
    "for epoch in range(hyper_params[\"num_epochs\"]):\n",
    "    trn = []\n",
    "    net.train()\n",
    "    for i, (images, labels) in enumerate(data.train_dl) :\n",
    "        if torch.cuda.is_available():\n",
    "            images = torch.autograd.Variable(images).cuda().float()\n",
    "            labels = torch.autograd.Variable(labels).cuda()\n",
    "        else : \n",
    "            images = torch.autograd.Variable(images).float()\n",
    "            labels = torch.autograd.Variable(labels)\n",
    "\n",
    "        y_pred = net(images)\n",
    "\n",
    "        loss = F.cross_entropy(y_pred, labels)\n",
    "        trn.append(loss.item())\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "#         torch.nn.utils.clip_grad_value_(net.parameters(), 10)\n",
    "        optimizer.step()\n",
    "\n",
    "        if i % 50 == 49 :\n",
    "            print('epoch = ', epoch, ' step = ', i + 1, ' of total steps ', total_step, ' loss = ', loss.item())\n",
    "\n",
    "    train_loss = (sum(trn) / len(trn))\n",
    "    train_loss_list.append(train_loss)\n",
    "\n",
    "    net.eval()\n",
    "    val = []\n",
    "    with torch.no_grad() :\n",
    "        for i, (images, labels) in enumerate(data.valid_dl) :\n",
    "            if torch.cuda.is_available():\n",
    "                images = torch.autograd.Variable(images).cuda().float()\n",
    "                labels = torch.autograd.Variable(labels).cuda()\n",
    "            else : \n",
    "                images = torch.autograd.Variable(images).float()\n",
    "                labels = torch.autograd.Variable(labels)\n",
    "\n",
    "            # Forward pass\n",
    "            y_pred = net(images)\n",
    "            \n",
    "            loss = F.cross_entropy(y_pred, labels)\n",
    "            val.append(loss.item())\n",
    "\n",
    "    val_loss = sum(val) / len(val)\n",
    "    val_loss_list.append(val_loss)\n",
    "    val_acc = _get_accuracy(data.valid_dl, net)\n",
    "\n",
    "    print('epoch : ', epoch + 1, ' / ', hyper_params[\"num_epochs\"], ' | TL : ', train_loss, ' | VL : ', val_loss, ' | VA : ', val_acc * 100)\n",
    "\n",
    "    if (val_acc * 100) > min_val :\n",
    "        print('saving model')\n",
    "        min_val = val_acc * 100\n",
    "        torch.save(net.state_dict(), '../saved_models/small_classifier/model1.pt')savename"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "        <style>\n",
       "            /* Turns off some styling */\n",
       "            progress {\n",
       "                /* gets rid of default border in Firefox and Opera. */\n",
       "                border: none;\n",
       "                /* Needs to be in here for Safari polyfill so background images work as expected. */\n",
       "                background-size: auto;\n",
       "            }\n",
       "            .progress-bar-interrupted, .progress-bar-interrupted::-webkit-progress-bar {\n",
       "                background: #F44336;\n",
       "            }\n",
       "        </style>\n",
       "      <progress value='142' class='' max='151', style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      94.04% [142/151 01:07<00:04]\n",
       "    </div>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "tfms = get_transforms(do_flip=False)\n",
    "data = ImageDataBunch.from_folder(new_path, train = 'train', valid = 'test', bs = hyper_params[\"batch_size\"], size = 224, ds_tfms = tfms).normalize(imagenet_stats)\n",
    "\n",
    "learn = cnn_learner(data, models.resnet34, metrics = accuracy)\n",
    "learn = learn.load('/home/akshay/.fastai/data/imagenette/models/resnet34_imagenette_bs64')\n",
    "learn.freeze()\n",
    "# learn.summary()\n",
    "\n",
    "class Flatten(nn.Module) :\n",
    "    def forward(self, input):\n",
    "        return input.view(input.size(0), -1)\n",
    "\n",
    "def conv2(ni, nf) : \n",
    "    return conv_layer(ni, nf, stride = 2)\n",
    "\n",
    "class ResBlock(nn.Module):\n",
    "    def __init__(self, nf):\n",
    "        super().__init__()\n",
    "        self.conv1 = conv_layer(nf,nf)\n",
    "        \n",
    "    def forward(self, x): \n",
    "        return (x + self.conv1(x))\n",
    "\n",
    "def conv_and_res(ni, nf): \n",
    "    return nn.Sequential(conv2(ni, nf), ResBlock(nf))\n",
    "\n",
    "def conv_(nf) : \n",
    "    return nn.Sequential(conv_layer(nf, nf), ResBlock(nf))\n",
    "    \n",
    "net = nn.Sequential(\n",
    "    conv_layer(3, 64, ks = 7, stride = 2, padding = 3),\n",
    "    nn.MaxPool2d(3, 2, padding = 1),\n",
    "    conv_(64),\n",
    "    conv_and_res(64, 128),\n",
    "    conv_and_res(128, 256),\n",
    "    conv_and_res(256, 512),\n",
    "    AdaptiveConcatPool2d(),\n",
    "    Flatten(),\n",
    "    nn.Linear(2 * 512, 256),\n",
    "    nn.Linear(256, 10)\n",
    ")\n",
    "if torch.cuda.is_available() : \n",
    "    net.cuda()\n",
    "\n",
    "preds, y, losses = learn.get_preds(with_loss=True, ds_type = DatasetType.Valid)\n",
    "interp = ClassificationInterpretation(learn, preds, y, losses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "99.99466666666666\n",
      "0.9999466666666667\n"
     ]
    }
   ],
   "source": [
    "conf_mat = interp.confusion_matrix()\n",
    "acc = np.trace(conf_mat) / np.sum(conf_mat, axis = None)\n",
    "print(acc * 100)\n",
    "\n",
    "print(_get_accuracy(data.valid_dl, learn.model))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "79.10133333333333\n",
      "99.99466666666666\n",
      "62.114666666666665\n"
     ]
    }
   ],
   "source": [
    "net.cpu()\n",
    "net.load_state_dict(torch.load('../saved_models/imagenette/less_data_classifier/model0.pt', map_location = 'cpu'))\n",
    "net.cuda()\n",
    "\n",
    "learn = cnn_learner(data, models.resnet34, metrics = accuracy)\n",
    "learn = learn.load('/home/akshay/.fastai/data/imagenette/models/resnet34_imagenette_bs64')\n",
    "learn.freeze()\n",
    "\n",
    "print(_get_accuracy(data.valid_dl, net) * 100)\n",
    "print(_get_accuracy(data.valid_dl, learn.model) * 100)\n",
    "\n",
    "net.cpu()\n",
    "net.load_state_dict(torch.load('../saved_models/imagenette/less_data_medium_no_teacher/model0.pt', map_location = 'cpu'))\n",
    "net.cuda()\n",
    "print(_get_accuracy(data.valid_dl, net) * 100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "from fastai.vision import *\n",
    "import torch\n",
    "from torchsummary import summary\n",
    "torch.cuda.set_device(0)\n",
    "\n",
    "for repeated in range(0, 1) : \n",
    "    torch.manual_seed(repeated)\n",
    "    torch.cuda.manual_seed(repeated)\n",
    "\n",
    "    # stage should be in 0 to 5 (5 for classifier stage)\n",
    "    hyper_params = {\n",
    "        \"stage\": 5,\n",
    "        \"repeated\": repeated,\n",
    "        \"num_classes\": 10,\n",
    "        \"batch_size\": 64,\n",
    "        \"num_epochs\": 100,\n",
    "        \"learning_rate\": 1e-4\n",
    "    }\n",
    "\n",
    "    path = untar_data(URLs.IMAGENETTE)\n",
    "    tfms = get_transforms(do_flip=False)\n",
    "    data = ImageDataBunch.from_folder(path, train = 'train', valid = 'val', bs = hyper_params[\"batch_size\"], size = 224, ds_tfms = tfms).normalize(imagenet_stats)\n",
    "    \n",
    "    class Flatten(nn.Module) :\n",
    "        def forward(self, input):\n",
    "            return input.view(input.size(0), -1)\n",
    "\n",
    "    def conv2(ni, nf) : \n",
    "        return conv_layer(ni, nf, stride = 2)\n",
    "\n",
    "    class ResBlock(nn.Module):\n",
    "        def __init__(self, nf):\n",
    "            super().__init__()\n",
    "            self.conv1 = conv_layer(nf,nf)\n",
    "\n",
    "        def forward(self, x): \n",
    "            return (x + self.conv1(x))\n",
    "\n",
    "    def conv_and_res(ni, nf): \n",
    "        return nn.Sequential(conv2(ni, nf), ResBlock(nf))\n",
    "\n",
    "    def conv_(nf) : \n",
    "        return nn.Sequential(conv_layer(nf, nf), ResBlock(nf))\n",
    "\n",
    "    net = nn.Sequential(\n",
    "        conv_layer(3, 64, ks = 7, stride = 2, padding = 3),\n",
    "        nn.MaxPool2d(3, 2, padding = 1),\n",
    "        conv_(64),\n",
    "        conv_and_res(64, 128),\n",
    "        conv_and_res(128, 256),\n",
    "        AdaptiveConcatPool2d(),\n",
    "        Flatten(),\n",
    "        nn.Linear(2 * 256, 128),\n",
    "        nn.Linear(128, hyper_params[\"num_classes\"])\n",
    "    )\n",
    "\n",
    "    net.cpu()\n",
    "    filename = '../saved_models/small_stage4/model' + str(repeated) + '.pt'\n",
    "    net.load_state_dict(torch.load(filename, map_location = 'cpu'))\n",
    "\n",
    "    if torch.cuda.is_available() : \n",
    "        net = net.cuda()\n",
    "        print('Model on GPU')\n",
    "\n",
    "    for name, param in net.named_parameters() : \n",
    "        param.requires_grad = False\n",
    "        if name[0] == '7' or name[0] == '8':\n",
    "            param.requires_grad = True\n",
    "        \n",
    "    optimizer = torch.optim.Adam(net.parameters(), lr = hyper_params[\"learning_rate\"])\n",
    "    total_step = len(data.train_ds) // hyper_params[\"batch_size\"]\n",
    "    train_loss_list = list()\n",
    "    val_loss_list = list()\n",
    "    min_val = 0\n",
    "    savename = '../saved_models/small_classifier/model' + str(repeated) + '.pt'\n",
    "    for epoch in range(hyper_params[\"num_epochs\"]):\n",
    "        trn = []\n",
    "        net.train()\n",
    "        for i, (images, labels) in enumerate(data.train_dl) :\n",
    "            if torch.cuda.is_available():\n",
    "                images = torch.autograd.Variable(images).cuda().float()\n",
    "                labels = torch.autograd.Variable(labels).cuda()\n",
    "            else : \n",
    "                images = torch.autograd.Variable(images).float()\n",
    "                labels = torch.autograd.Variable(labels)\n",
    "\n",
    "            y_pred = net(images)\n",
    "\n",
    "            loss = F.cross_entropy(y_pred, labels)\n",
    "            trn.append(loss.item())\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "    #         torch.nn.utils.clip_grad_value_(net.parameters(), 10)\n",
    "            optimizer.step()\n",
    "\n",
    "            if i % 50 == 49 :\n",
    "                print('epoch = ', epoch, ' step = ', i + 1, ' of total steps ', total_step, ' loss = ', loss.item())\n",
    "\n",
    "        train_loss = (sum(trn) / len(trn))\n",
    "        train_loss_list.append(train_loss)\n",
    "\n",
    "        net.eval()\n",
    "        val = []\n",
    "        with torch.no_grad() :\n",
    "            for i, (images, labels) in enumerate(data.valid_dl) :\n",
    "                if torch.cuda.is_available():\n",
    "                    images = torch.autograd.Variable(images).cuda().float()\n",
    "                    labels = torch.autograd.Variable(labels).cuda()\n",
    "                else : \n",
    "                    images = torch.autograd.Variable(images).float()\n",
    "                    labels = torch.autograd.Variable(labels)\n",
    "\n",
    "                # Forward pass\n",
    "                y_pred = net(images)\n",
    "\n",
    "                loss = F.cross_entropy(y_pred, labels)\n",
    "                val.append(loss.item())\n",
    "\n",
    "        val_loss = sum(val) / len(val)\n",
    "        val_loss_list.append(val_loss)\n",
    "        val_acc = _get_accuracy(data.valid_dl, net)\n",
    "\n",
    "        print('epoch : ', epoch + 1, ' / ', hyper_params[\"num_epochs\"], ' | TL : ', train_loss, ' | VL : ', val_loss, ' | VA : ', val_acc * 100)\n",
    "\n",
    "        if (val_acc * 100) > min_val :\n",
    "            print('saving model')\n",
    "            min_val = val_acc * 100\n",
    "            torch.save(net.state_dict(), savename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (ak_fastai)",
   "language": "python",
   "name": "ak_fastai"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
